#include <map>
#include <cassert>
#include <cstddef>
#include <cmath>
#include "reflection_pad1d_tiling.h"

/**
 *                       _oo0oo_
 *                      o8888888o
 *                      88" . "88
 *                      (| -_- |)
 *                      0\  =  /0
 *                    ___/`---'\___
 *                  .' \\|     |// '.
 *                 / \\|||  :  |||// \
 *                / _||||| -:- |||||- \
 *               |   | \\\  -  /// |   |
 *               | \_|  ''\---/''  |_/ |
 *               \  .-\__  '-'  ___/-. /
 *             ___'. .'  /--.--\  `. .'___
 *          ."" '<  `.___\_<<|>_/___.' >' "".
 *         | | :  `- \`.;`\ _ /`;.`/ - ` : | |
 *         \  \ `_.   \_ __\ /__ _/   .-` /  /
 *     =====`-.____`.___ \_____/___.-`___.-'=====
 *                       `=---='
 *
 *     ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 *               佛祖保佑         永无BUG
 */

constexpr int32_t X_INPUT_INDEX = 0;                  // 输入张量索引
constexpr int32_t PAD_INPUT_INDEX = 1;                // 填充参数索引
constexpr int32_t OUTPUT_INDEX = 0;                   // 输出张量索引

// 维度索引映射
constexpr uint32_t DIM_1D_W = 0;                      // 1D: [W]
constexpr uint32_t DIM_2D_C = 0;                      // 2D: [C, W]
constexpr uint32_t DIM_2D_W = 1;
constexpr uint32_t DIM_3D_N = 0;                      // 3D: [N, C, W]
constexpr uint32_t DIM_3D_C = 1;
constexpr uint32_t DIM_3D_W = 2;

// 填充参数索引
constexpr uint32_t PAD_LEFT_INDEX = 0;                // 左填充
constexpr uint32_t PAD_RIGHT_INDEX = 1;               // 右填充
constexpr uint32_t PAD_LEN = 2;                       // 填充参数长度

constexpr uint32_t BYTE_BLOCK = 32;                   // 32字节对齐
constexpr uint32_t MAX_CORE_NUM = 64;                 // 最大核心数
constexpr uint32_t BUFFER_NUM = 2;

// 向上对齐函数
template <typename T1, typename T2>
T1 CeilAlign(T1 value, T2 align) {
    if (align == 0) return value;
    return (value + align - 1) / align * align;
}

// 向下对齐函数
template <typename T1, typename T2>
T1 FloorAlign(T1 value, T2 align) {
    if (align == 0) return value;
    return value / align * align;
}

void GenerateTilingData(uint8_t *tilingBuf, uint32_t coreNum)
{
    assert(tilingBuf != nullptr && "tilingBuf must not be null (invalid input buffer)");
    assert(coreNum > 0 && "coreNum must be greater than 0 (invalid core count)");

    ReflectionPad1dTilingData *tiling = reinterpret_cast<ReflectionPad1dTilingData *>(tilingBuf);

    assert(tiling->nSize > 0 && "nSize must be > 0 (invalid batch dimension)");
    assert(tiling->cSize > 0 && "cSize must be > 0 (invalid channel dimension)");
    assert(tiling->wSize > 0 && "wSize must be > 0 (invalid width dimension)");
    assert(tiling->padLeft >= 0 && "padLeft must be non-negative (invalid padding)");
    assert(tiling->padRight >= 0 && "padRight must be non-negative (invalid padding)");

    uint32_t dtypeBytes;
    switch (tiling->tilingKey) {
        case static_cast<uint32_t>(TilingKey::TILING_KEY_FP32):
            dtypeBytes = FLOAT_BYTES;
            break;
        case static_cast<uint32_t>(TilingKey::TILING_KEY_FP16):
        case static_cast<uint32_t>(TilingKey::TILING_KEY_BF16):
            dtypeBytes = FLOAT16_BYTES;
            break;
        default:
            assert(false && "Invalid tilingKey (unsupported data type)");
            dtypeBytes = FLOAT_BYTES; 
    }

    uint32_t outWSize = tiling->wSize + tiling->padLeft + tiling->padRight;
    uint32_t elePer32B = BYTE_BLOCK / dtypeBytes;
    uint32_t alignWSize = CeilAlign(tiling->wSize, elePer32B);
    uint32_t alignOutWSize = CeilAlign(outWSize, elePer32B);

    uint64_t totalNc = static_cast<uint64_t>(tiling->nSize) * tiling->cSize;
    uint32_t blockNum = 0, ncPerCore = 0, tailNC = 0;

    if (totalNc <= coreNum) {
        // 任务量少：1个任务对应1个核心，剩余核心空闲
        blockNum = static_cast<uint32_t>(totalNc);
        ncPerCore = 1;
        tailNC = 0;
    } else {
        // 任务量多：平均分配，前tailNC个核心多处理1个任务
        ncPerCore = static_cast<uint32_t>(totalNc / coreNum);
        tailNC = static_cast<uint32_t>(totalNc % coreNum);
        blockNum = coreNum;
    }

    tiling->alignWSize = alignWSize;
    tiling->alignOutWSize = alignOutWSize;
    tiling->blockNum = blockNum;
    tiling->ncPerCore = ncPerCore;
    tiling->tailNC = tailNC;
}